/*
 * topological-hex, a program to compute combinatorial hexahedral meshes.
 *
 * Copyright (C) <2018> <Université catholique de Louvain (UCL), Belgique>
 *
 * List of the contributors to the development, description and complete
 * License: see LICENSE file.
 *
 * This program (topological-hex) is free software:
 * you can redistribute it and/or modify it under the terms
 * of the GNU General Public License as published by the Free
 * Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program (see COPYING file).  If not,
 * see <http://www.gnu.org/licenses/>.
 */

#include <check.h>
#include <pthread.h>
#include <unistd.h>
#include <stdlib.h>
#include <stdatomic.h>
#include <omp.h>

#include "input.h"

#include "server.h"

#define NUM_CLIENTS 4
#define NUM_THREADS_PER_CLIENT 2

#define MAX_HEX 64

typedef struct server_args {
  bool ready;
  atomic_uint_fast32_t num_solutions[MAX_HEX + 1];
} server_args;

typedef struct client_args {
  const char *filename;
  const char *symmetries;
  uint32_t max_hex;
  uint32_t num_internal_vertices;
} client_args;

static void on_solution(const char *name,
                        const vertex_index *vertices, hex_index num_hexes,
                        void *data) {
  atomic_fetch_add(((server_args*)data)->num_solutions + num_hexes, 1);
}

static void on_log(const char *message, void *data) {}

static void on_listen(void *data) {
  ((server_args*)data)->ready = true;
}

static void on_progress(uint32_t i, uint32_t n, void *data) {}

static void *server_fn(void *data) {
  server_args *args = data;
  server_run(on_solution, args, on_log, NULL, on_listen, args,
             on_progress, NULL,
             NULL);
  return NULL;
}

static void *master_client_fn(void *data) {
  client_args *args = data;

  mesh mesh;
  ck_assert(mesh_load(&mesh, args->filename) == SUCCESS);
  ck_assert(mesh_reserve(&mesh, args->max_hex) == SUCCESS);

  vertex_index *symmetries;
  vertex_index num_boundary_vertices;
  size_t num_symmetries;
  if (args->symmetries) {
    ck_assert(read_symmetries(args->symmetries,
                              &symmetries, &num_boundary_vertices,
                              &num_symmetries) == SUCCESS);
  }

  solver solver;
  ck_assert(solver_init(&solver, &mesh, mesh.num_vertices +
                        args->num_internal_vertices,
                        SOLVER_BIPARTITE_GRAPH) == SUCCESS);

  if (args->symmetries) {
    ck_assert(solver_set_symmetries(&solver, num_symmetries,
                                    num_boundary_vertices, symmetries) ==
              SUCCESS);
  }

  solver_run_distributed(&solver, "localhost",
                         on_log, NULL, 3, 256);

  solver_release(&solver);
  if (args->symmetries) free(symmetries);
  mesh_release(&mesh);

  return NULL;
}

static void *client_fn(void *args) {
  solver_run_distributed_client("localhost", on_log, NULL, 3);
  return NULL;
}

static void test_distributed(const char *input_name,
                             const char *symmetries,
                             uint32_t max_hex,
                             uint32_t num_internal_vertices,
                             const size_t *expected_sizes) {

  omp_set_num_threads(NUM_THREADS_PER_CLIENT);
  pthread_t server_thread, clients[NUM_CLIENTS];

  server_args args;
  args.ready = false;
  for (size_t i = 0; i <= MAX_HEX; i++)
    atomic_init(args.num_solutions + i, 0);

  ck_assert(pthread_create(&server_thread, NULL, server_fn, &args) == 0);

  while (!args.ready)
    usleep(1000);

  client_args client_args = {
    input_name, symmetries, max_hex, num_internal_vertices
  };

  ck_assert(pthread_create(&clients[0], NULL, master_client_fn, &client_args) == 0);
  for (size_t i = 1; i < NUM_CLIENTS; i++)
    ck_assert(pthread_create(&clients[i], NULL, client_fn, NULL) == 0);

  ck_assert(pthread_join(server_thread, NULL) == 0);
  for (size_t i = 0; i < NUM_CLIENTS; i++)
    ck_assert(pthread_join(clients[i], NULL) == 0);

  for (uint32_t i = 0; i <= MAX_HEX; i++) {
    ck_assert_msg(atomic_load(args.num_solutions + i) == expected_sizes[i],
                  "Got %u solution of size %u, expected %u",
                  atomic_load(args.num_solutions + i), i, expected_sizes[i]);
  }
}

START_TEST(test_cube) {
  size_t expected_sizes[MAX_HEX+1] = {0};
  expected_sizes[1]  = 1;
  expected_sizes[7] = 1;

  test_distributed(INPUT_DIR "/single-hex.mesh",
                   NULL,
                   MAX_HEX, 15,
                   expected_sizes);

  test_distributed(INPUT_DIR "/single-hex.mesh",
                   INPUT_DIR "/single-hex.symmetries",
                   MAX_HEX, 15,
                   expected_sizes);
} END_TEST

START_TEST(test_prism) {
  size_t expected_sizes[MAX_HEX+1] = {0};
  expected_sizes[6] = 1;

  test_distributed(INPUT_DIR "/prism.mesh",
                   NULL,
                   9, 1, expected_sizes);
} END_TEST

START_TEST(test_rhombic_dodecahedron) {
  size_t expected_sizes[MAX_HEX+1] = {0};
  expected_sizes[4]  = 2;

  test_distributed(INPUT_DIR "/rhombic-dodecahedron.mesh",
                   NULL,
                   MAX_HEX, 6,
                   expected_sizes);
} END_TEST

START_TEST(test_rhombic_dodecahedron_symmetric) {
  size_t expected_sizes[MAX_HEX+1] = {0};
  expected_sizes[4]  = 1;
  expected_sizes[10] = 1;

  test_distributed(INPUT_DIR "/rhombic-dodecahedron.mesh",
                   INPUT_DIR "/rhombic-dodecahedron.symmetries",
                   MAX_HEX, 10,
                   expected_sizes);
} END_TEST

Suite *distributed_suite(void) {
  Suite *s = suite_create("Distributed Execution");

  TCase *tcase = tcase_create("Distributed Execution");
  tcase_set_timeout(tcase, 15.0);
  tcase_add_test(tcase, test_cube);
  tcase_add_test(tcase, test_prism);
  tcase_add_test(tcase, test_rhombic_dodecahedron);
  tcase_add_test(tcase, test_rhombic_dodecahedron_symmetric);
  suite_add_tcase(s, tcase);

  return s;
}
